# Copyright 2025 The Lynx Authors. All rights reserved.
# Licensed under the Apache License Version 2.0 that can be found in the
# LICENSE file in the root directory of this source tree.

#!/usr/bin/python
# -*- coding: UTF-8 -*-


import os
import sys
import time
import yaml
from jni_generator import GenerateJNIHeader, Options

def parse_yaml(input_file_path):
  with open(input_file_path, 'r') as file:
    data = yaml.safe_load(file)
    return data

so_load_file_template = """
// This file is autogenerated.

CUSTOM_HEADERS
// AUTO_GENERATED_INCLUDE_HEADERS_START
// AUTO_GENERATED_INCLUDE_HEADERS_END

NAMESPACE_START
extern "C" JNIEXPORT jint JNI_OnLoad(JavaVM* vm, void* reserved) {
  lynx::base::android::InitVM(vm);
  JNIEnv* env = base::android::AttachCurrentThread();
  // AUTO_GENERATED_REGISTER_METHODS_START
  // AUTO_GENERATED_REGISTER_METHODS_END
  return JNI_VERSION_1_6;
}

NAMESPACE_END
"""

build_gn_template = """
# This file is autogenerated.

CUSTOM_HEADERS
TEMPLATE_NAME("build") {
  sources = [
    # AUTO_GENERATED_FILES_START
    # AUTO_GENERATED_FILES_END
  ]
}
"""

register_header_template = """
// This file is autogenerated for
//     JAVA_FILE_PATH

#ifndef HEADER_GUARD
#define HEADER_GUARD

#include <jni.h>

namespace lynx {
namespace jni {

bool FUNCTION_NAME(JNIEnv* env);

}  // namespace jni
}  // namespace lynx

#endif  // HEADER_GUARD

"""

def generate_register_file(java_file, java_base_name, function_name, gen_path, output_path):
  # generate XXX_register_jni.h which contains RegisterJNIForXXX method.
  guard_string = java_file.replace('/', '_')
  guard_string = guard_string.split('.')[0]
  guard_string = guard_string + '_REGISTER_JNI_H'
  header_filled_template = register_header_template.replace('HEADER_GUARD', guard_string)
  header_filled_template = header_filled_template.replace('JAVA_FILE_PATH', java_file)
  header_filled_template = header_filled_template.replace('FUNCTION_NAME', function_name)
  header_file = os.path.join(output_path, java_base_name + '_register_jni.h')
  with open(header_file, 'w') as file:
    file.write(header_filled_template)

def append_content_if_changed(file_path, start_flag, end_flag, new_content_list):
  # Find content that starts at start_flag and ends at end_flag
  file_lines = []
  matched_lines = []
  start = False
  end = False
  file_start_num = 0
  file_end_num = 0
  with open(file_path, 'r') as file:
    index = 0
    while True:
        line = file.readline()
        if not line:
          break
        file_lines.append(line)
        index = index + 1
        if not start and start_flag in line:
          start = True
          file_start_num = index
          continue
        if not end and end_flag in line:
          end = True
          file_end_num = index
        if (not start) or (start and end):
          continue
        matched_lines.append(line)

  # Find where the new content begins in old content.
  found = False
  first_line = new_content_list[0]
  start_line_num = 0
  end_line_num = 0
  for index, line in enumerate(matched_lines):
    if first_line in line:
      found = True
      start_line_num = index
      
    if found and line == '\n':
      end_line_num = index
      break

  # Insert new content to file
  changed = False
  if not found:
    changed = True
  elif found and end_line_num - start_line_num != len(new_content_list):
    changed = True
    # Remove old content
    start = file_start_num + start_line_num
    end = file_start_num + end_line_num + 1
    print('delete line {} to {}'.format(start, end))
    del file_lines[start:end]
  else:
    for index, line in enumerate(new_content_list):
      if not line == matched_lines[index + start_line_num]:
        changed = True
        # Remove old content
        start = file_start_num + start_line_num
        end = file_start_num + end_line_num + 1
        print('delete line {} to {}'.format(start, end))
        del file_lines[start:end]
        break
  
  if changed:
    # Write new content to file
    print('file changed')
    new_content_string = ''.join(new_content_list)
    file_lines.insert(file_start_num, new_content_string + '\n')
    with open(file_path, 'w') as file:
      file.writelines(file_lines)

# Write content to XXXSoLoad.cpp
def append_content_to_so_registry(so_configs, include_headers, register_methods):
  cpp_output_path = so_configs.get('output_path', '')
  if not os.path.exists(cpp_output_path):
    # If file not exist, write template content to XXXSoLoad.cpp
    custom_headers = so_configs.get('custom_headers', [])
    namespaces = so_configs.get('namespaces', ['lynx'])
    so_file_str = so_load_file_template
    if len(custom_headers) > 0:
      header_str = ''
      for custom_header in custom_headers:
        header_str = header_str + '#include "{}"\n'.format(custom_header)
      so_file_str = so_file_str.replace('CUSTOM_HEADERS', header_str)
    
    namespace_start_str = ''
    namespace_end_str = ''
    for namespace in namespaces:
      namespace_start_str =  namespace_start_str + 'namespace {} '.format(namespace) + '{\n'
      namespace_end_str = '}' + '  // namespace {}\n'.format(namespace) + namespace_end_str
    so_file_str = so_file_str.replace('NAMESPACE_START', namespace_start_str)
    so_file_str = so_file_str.replace('NAMESPACE_END', namespace_end_str)

    directory = os.path.dirname(cpp_output_path)
    if not os.path.exists(directory):
      os.makedirs(directory)
    with open(cpp_output_path, 'w') as file:
      file.write(so_file_str)

  # Write include headers to XXXSoLoad.cpp
  include_headers_list = []
  for include_header in include_headers:
    if len(include_header[1]) != 0:
      include_headers_list.append('#if ' + include_header[1] + '\n')
      include_headers_list.append(include_header[0] + '\n')
      include_headers_list.append('#endif\n')
    else:
      include_headers_list.append(include_header[0] + '\n')
  append_content_if_changed(cpp_output_path, 
                            'AUTO_GENERATED_INCLUDE_HEADERS_START',
                            'AUTO_GENERATED_INCLUDE_HEADERS_END',
                            include_headers_list)

  # Write register method to XXXSoLoad.cpp
  register_methods_list = []
  for register_method in register_methods:
    new_line = '  ' + register_method[0] + '\n'
    if len(register_method[1]) != 0:
      register_methods_list.append('#if ' + register_method[1] + '\n')
      register_methods_list.append(new_line)
      register_methods_list.append('#endif\n')
    else:
      register_methods_list.append(new_line)

  append_content_if_changed(cpp_output_path, 
                            'AUTO_GENERATED_REGISTER_METHODS_START',
                            'AUTO_GENERATED_REGISTER_METHODS_END',
                            register_methods_list)
  
# Write content to BUILD.gn
def append_files_to_gn(gn_configs, gn_files):
  gn_file_path = gn_configs.get('gn_file_path', '')

  if not os.path.exists(gn_file_path):
    # If file not exist, write template content to BUILD.gn.
    custom_headers = gn_configs.get('custom_headers', [])
    template_name = gn_configs.get('template_name', 'source_set')
    gn_file_str = build_gn_template
    if len(custom_headers) > 0:
      header_str = ''
      for custom_header in custom_headers:
        header_str = header_str + 'import("{}")\n'.format(custom_header)
      gn_file_str = gn_file_str.replace('CUSTOM_HEADERS', header_str)
    gn_file_str = gn_file_str.replace('TEMPLATE_NAME', template_name)

    with open(gn_file_path, 'w') as file:
      file.write(gn_file_str)

  # Write new content to BUILD.gn
  gn_sources_list = []
  for gn_file in gn_files:
    new_line = '    ' + gn_file + '\n'
    gn_sources_list.append(new_line)
  append_content_if_changed(gn_file_path, 
                            'AUTO_GENERATED_FILES_START',
                            'AUTO_GENERATED_FILES_END',
                            gn_sources_list)
  

def generate_files(root_path, java_root_path, jni_output_path, include_root_path):
  # Parse jni_files yaml file to a map
  parent_path = os.path.dirname(jni_output_path)
  input_file_path = os.path.join(root_path, parent_path, 'jni_files.yml')
  jni_configs = parse_yaml(input_file_path)

  # Read config from yaml map
  inputs = jni_configs.get('inputs', [])
  special_cases = jni_configs.get('special_cases', [])
  so_configs = jni_configs.get('so_load_configs', {})
  cpp_output_path = so_configs.get('output_path', '')
  gn_configs = jni_configs.get('gn_configs', {})

  # Handle special cases that are shared or not automatically generated
  excluded_java_files = {}
  special_headers = []
  special_methods = []
  for special_case in special_cases:
    header = special_case.get('header', '')
    method = special_case.get('method', '')
    java = special_case.get('java', '')
    macro = special_case.get('macro', '')
    if header != '':
      special_headers.append(['#include "' + header + '"', macro])
    if method != '':
      special_methods.append([method, macro])
    if java != '':
      excluded_java_files[java] = True

  # Read Java files and assemble include header and register method and gn files
  # e.g.:
  #  com/lynx/tasm/behavior/PaintingContext.java
  #    include header is #include "${path}/PaintingContext_register_jni.h"
  #    register method is RegisterJNIForPaintingContext(env)
  hash_map = {}
  include_headers = []
  register_methods = []
  gn_files = []
  for input in inputs:
    java_path = input.get('java', '')
    method_name = input.get('method_name', '')
    macro = input.get('macro', '')
    if len(java_path) == 0:
      continue
    if hash_map.get(java_path):
      continue
    if excluded_java_files.get(java_path):
      continue
    java_file_full_path = os.path.join(java_root_path, java_path)
    if not os.path.exists(java_file_full_path):
      continue
    hash_map[java_path] = True
    java_file_name = os.path.basename(java_path)
    java_base_name = os.path.splitext(java_file_name)[0]
    if len(method_name) == 0:
      method_name = 'RegisterJNIFor' + java_base_name
    output_jni_file_name = java_base_name + '_jni.h'
    output_full_path = os.path.join(root_path, jni_output_path, output_jni_file_name)
    include_header = '#include "' + include_root_path + '/' + java_base_name + '_register_jni.h"'
    include_headers.append([include_header, macro])
    register_method = 'lynx::jni::' + method_name + '(env);'
    register_methods.append([register_method, macro])
    gn_files.append('"gen/' + output_jni_file_name + '",')
    gn_files.append('"gen/' + java_base_name + '_register_jni.h",')

    # generate jni
    options = Options(False)
    print(output_full_path)
    GenerateJNIHeader(java_file_full_path, output_full_path, options)
    # generate register method file
    generate_register_file(java_path, 
                           java_base_name, 
                           method_name, 
                           include_root_path, 
                           jni_output_path)
  # generate SoLoad.cpp
  include_headers.extend(special_headers)
  register_methods.extend(special_methods)
  include_headers.sort()
  register_methods.sort()
  cpp_full_path = os.path.join(root_path, cpp_output_path)
  so_configs['output_path'] = cpp_full_path
  append_content_to_so_registry(so_configs, include_headers, register_methods)
  # generate BUILD.gn
  gn_files.sort()
  gn_configs['gn_file_path'] = os.path.join(root_path, parent_path, 'BUILD.gn')
  append_files_to_gn(gn_configs, gn_files)

def main():
  if len(sys.argv) <= 4:
    print('The params count is less than 4\n')
    return -1
  generate_files(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])
  return 0

if __name__ == "__main__":
  sys.exit(main())
